#region Copyright
/// <copyright>
/// Copyright (c) 2007 Ilia Shramko.
/// Released under the terms of the GNU General Public License.
/// </copyright>
#endregion

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Text;

namespace Osh.Dvb
{
	public class TransportStreamDemultiplexer
	{
		[Flags]
		private enum AdaptationFieldControl
		{
			Payload = 1,
			AdaptationField = 2,
		}

		#region StreamInfo Class
		private class StreamInfo
		{
			private List<IDvbPacketStreamFilter> _tsFilters;
			private List<IDvbPacketStreamFilter> _typedFilters;
			private PacketStreamType _type;
			private sbyte _continuityCounter;
			private byte[] _packetBuffer;
			private int _packetSize;

			public StreamInfo(PacketStreamType type)
			{
				_type = type;
				_packetBuffer = null;
				_packetSize = 0;
			}

			public void AddFilter(PacketStreamType type, IDvbPacketStreamFilter filter)
			{
				if(type == PacketStreamType.TS)
				{
					if(_tsFilters == null)
						_tsFilters = new List<IDvbPacketStreamFilter>();
					_tsFilters.Add(filter);
				}
				else
				{
					if(_type == PacketStreamType.TS)
						_type = type;
					else if(_type != type)
						throw new ArgumentOutOfRangeException("type");

					if(_typedFilters == null)
						_typedFilters = new List<IDvbPacketStreamFilter>();
					_typedFilters.Add(filter);
				}
			}

			public void RemoveFilter(IDvbPacketStreamFilter filter)
			{
				if(_tsFilters != null)
				{
					_tsFilters.Remove(filter);
					if(_tsFilters.Count == 0)
						_tsFilters = null;
				}
				if(_typedFilters != null)
				{
					_typedFilters.Remove(filter);
					if(_typedFilters.Count == 0)
					{
						_typedFilters = null;
						_type = PacketStreamType.TS;
					}
				}
			}

			public PacketStreamType Type
			{
				get { return _type; }
			}

			public List<IDvbPacketStreamFilter> TSFilters
			{
				get { return _tsFilters; }
			}

			public List<IDvbPacketStreamFilter> TypedFilters
			{
				get { return _typedFilters; }
			}

			public bool HasTSFilters
			{
				get { return _tsFilters != null && _tsFilters.Count > 0; }
			}

			public bool HasTypedFilters
			{
				get { return _typedFilters != null && _typedFilters.Count > 0; }
			}

			public sbyte ContinuityCounter
			{
				get { return _continuityCounter; }
				set { _continuityCounter = value; }
			}

			public byte[] PacketBuffer
			{
				get { return _packetBuffer; }
			}

			public int PacketSize
			{
				get { return _packetSize; }
				set { _packetSize = value; }
			}

			public void EnsurePacketSize(int size)
			{
				if(_packetBuffer == null)
				{
					_packetBuffer = new byte[size];
					return;
				}
					
				if(_packetBuffer.Length < size)
				{
					byte[] newPacketBuffer = new byte[size];
					Buffer.BlockCopy(_packetBuffer, 0, newPacketBuffer, 0, _packetSize);
					_packetBuffer = newPacketBuffer;
				}
			}
		}
		#endregion StreamInfo Class

		private class AdaptationField
		{
			public bool DiscontinuityIndicator; // discontinuity_indicator
		}

		public const int TransportStreamPacketSize = 188;
		public const int MaxPid = 8192; // 2 ^ 13;
		private const int TransportErrorIndicator = 0x80;	// transport_error_indicator
		private const int PayloadStartIndicator = 0x40;	// payload_unit_start_indicator

		private IPacketStream _transportStream;
		private List<StreamInfo> _streams;
		private AdaptationField _adaptationField;

		public TransportStreamDemultiplexer(IPacketStream transportStream)
		{
			_transportStream = transportStream;
			_streams = new List<StreamInfo>(MaxPid);
		}

		private StreamInfo GetStream(int pid)
		{
			if(pid < _streams.Count)
				return _streams[pid];
			return null;
		}

		private void SetStream(int pid, StreamInfo si)
		{
			if(pid >= _streams.Count)
			{
				for(int i = _streams.Count; i <= pid; i++)
					_streams.Add(null);
			}
			_streams[pid] = si;
		}

		public void AddFilter(int pid, PacketStreamType type, IDvbPacketStreamFilter filter)
		{
			if(pid > MaxPid)
				throw new ArgumentOutOfRangeException("pid");

			StreamInfo si = GetStream(pid);
			if(si == null)
			{
				si = new StreamInfo(type);
				SetStream(pid, si);
			}

			si.AddFilter(type, filter);
		}

		public void RemoveFilter(int pid, IDvbPacketStreamFilter filter)
		{
			if(pid > MaxPid)
				throw new ArgumentOutOfRangeException("pid");

			StreamInfo si = GetStream(pid);
			if(si != null)
				si.RemoveFilter(filter);
		}

		public PacketStream OpenStream(int pid, PacketStreamType type)
		{
			TransportSubstream stream = new TransportSubstream(this, pid, type);
			AddFilter(pid, type, stream);
			return stream;
		}

		public bool Read()
		{
			return Read(false);
		}

		private bool Read(bool readAll)
		{
			byte[] packet;
			do
			{
				if(!_transportStream.MoveToNextPacket())
					return false;

				packet = _transportStream.PacketBuffer;
			    if(packet[0] != 0x47 || _transportStream.PacketSize != TransportStreamPacketSize)	// sync_byte
			        continue;
				
			    int indicators = packet[1];
			    if((indicators & TransportErrorIndicator) == TransportErrorIndicator)
			        continue;

			    int pid = ((indicators & 0x1F) << 8) + packet[2];	// PID
				StreamInfo streamInfo = GetStream(pid);
			    if(streamInfo == null)
			        continue;

				if(streamInfo.HasTSFilters)
					NotifyFilters(streamInfo.TSFilters, pid, packet, 0, TransportStreamPacketSize);
				else if(!streamInfo.HasTypedFilters)
					continue;


				if((indicators & PayloadStartIndicator) == PayloadStartIndicator)
				{
					if(streamInfo.PacketSize > 0)
					{
						if(streamInfo.Type == PacketStreamType.PSI)
							FeedPsiPacket(streamInfo, pid);
						else if(streamInfo.Type == PacketStreamType.PES)
							FeedPesPacket(streamInfo, pid);
						else
							streamInfo.PacketSize = 0;
					}
				}
				else
				{
					if(streamInfo.PacketSize == 0)
						continue;	// skip the rest of the packet if we don't have the beginning
				}

			    // transport_scrambling_control (2bit), adaptation_field_control(2bit), continuity_counter(4bit)
			    byte control = packet[3];
			    int transportScramblingControl = control >> 6;
			    int adaptationFieldControl = (control >> 4) & 0x03;
			    sbyte continuityCounter = (sbyte)(control & 0x0F);

				int position = 4;
				if((adaptationFieldControl & (int)AdaptationFieldControl.AdaptationField) != 0)
					if(!ReadAdaptationField(packet, ref position))
						continue;

				if((adaptationFieldControl & (int)AdaptationFieldControl.Payload) != 0)
				{
					// Check continuityCounter
					if(streamInfo.ContinuityCounter == continuityCounter)
						continue;

					streamInfo.ContinuityCounter = continuityCounter;

					int packetPayloadSize = TransportStreamPacketSize - position;

					streamInfo.EnsurePacketSize(streamInfo.PacketSize + packetPayloadSize);
					Buffer.BlockCopy(packet, position, streamInfo.PacketBuffer, streamInfo.PacketSize, packetPayloadSize);
					streamInfo.PacketSize = streamInfo.PacketSize + packetPayloadSize;
				}

			}
			while(readAll);

			return true;
		}

		private bool ReadAdaptationField(byte[] packet, ref int position)
		{
			int length = packet[position++];

			int flags = packet[position];

			_adaptationField = new AdaptationField();
			_adaptationField.DiscontinuityIndicator = (flags >> 7) != 0;

			position += length;
			return true;
		}

		private void FeedPsiPacket(StreamInfo streamInfo, int pid)
		{
			if(streamInfo.PacketSize >= 4)
			{
				byte[] buffer = streamInfo.PacketBuffer;
				int packetStart = 1 + buffer[0];	// pointer_field
				int packetSize = 3 + ((buffer[2] & 0xF) << 8) + buffer[3]; // section_length
				if(packetStart + packetSize <= streamInfo.PacketSize)
					NotifyFilters(streamInfo.TypedFilters, pid, buffer, packetStart, packetSize);
			}
			streamInfo.PacketSize = 0;
		}

		private void FeedPesPacket(StreamInfo streamInfo, int pid)
		{
			if(streamInfo.PacketSize >= 4)
			{
				byte[] buffer = streamInfo.PacketBuffer;
				if(buffer[0] == 0 && buffer[1] == 0 && buffer[2] == 1) // check packet_start_code_prefix for 0x000001
				{
					int packetSize = 6 + (buffer[4] << 8) + buffer[5]; // PES_packet_length
					if(packetSize == 6 || packetSize <= streamInfo.PacketSize)
						NotifyFilters(streamInfo.TypedFilters, pid, buffer, 0, packetSize);
				}
			}
			streamInfo.PacketSize = 0;
		}

		private void NotifyFilters(IList<IDvbPacketStreamFilter> filters, int pid, byte[] packet, int position, int length)
		{
			int numberOfFilters = filters.Count;
			for(int i = 0; i < numberOfFilters; i++)
				filters[i].NextPacket(pid, packet, position, length);
		}
	}
}
